from pydantic import BaseModel, Field
import json
import traceback
import capture
from transformBasePos import TransformPos
from DepthColorProcessor import DepthColorProcessor
from Serial_Control import SerialControl
import tts
import capFaceRecog
from apscheduler.schedulers.background import BackgroundScheduler#type: ignore
import time
from legMotion import legMotion
from armMotion import armMotion
from LLMservice import LLMService
import math
from typing import Union
from wheelMotion import wheelMotion
scheduler = BackgroundScheduler()
scheduler.start()
from config import port,portWheel
serial=SerialControl(port,115200)
serialW=SerialControl(portWheel,9600)
dcp=DepthColorProcessor(interactive_plots=False)
transformPos=TransformPos(serial,dcp)
armmotion=armMotion(serial,dcp)
legmotion=legMotion(serial)
wheelmotion=wheelMotion(serialW) 
class recogObj(BaseModel):
    label:str=Field(...,description="the name of the object")
    label_id:int=Field(default=-1,description="Don't care about it.")
    pos:list[float] = Field(...,description="x,y coordinates of the object's center in image's pixel")
class registerRecogResult(BaseModel):
    objs:list[recogObj]=Field(...,description="register the recog result that you recognized")
def Forward(distance: int):
    # for _ in range(distance // 10):  # Convert cm to steps
    #     legmotion.walk()
    # return {"success":True}
    pass
def Turn(direction: int):
    # steps = abs(direction) // 45  # Convert degrees to steps
    # if direction < 0:
    #     for _ in range(steps):
    #         legmotion.leftround()
    # else:
    #     for _ in range(steps):
    #         legmotion.rightround()
    # return {"success":True}
    pass
def ArmMove(x: float, y: float, z: float, releaseOnPos: bool, catchOnPos: bool):
    armmotion.toPos(x, y, z)
    if releaseOnPos:
        armmotion.release()
    elif catchOnPos:
        armmotion.catch()
    return {"success":True}
def Capture(yaw_degree: float, pitch_degree: float):
    armmotion.toDir(yaw_degree,pitch_degree)
    base64str = capture.fromImg("yolo.png")
    return {"success":True,"imagestr":base64str}
def Speak(content: str):
    print(f"speaking: {content}")
    tts.main(content)
    return {"success":True}
def SetSchedule(jobtype: str, task: str, hour: int, minute: int, second: int, taskagent):
    def job():
        taskagent.main(json.loads(task))
    if jobtype=="cron":
        scheduler.add_job(job, 'cron', hour=hour, minute=minute,second=second)
    elif jobtype=="interval":
        scheduler.add_job(job, 'interval',hours=hour, minutes=minute, seconds=second)
    else:
        return {"error": "Invalid job type, choose from cron or interval"}
    return {"success":True}
def yolo_json(filter_label_ids):
    time.sleep(1)
    with open("detected_objects.json", "r") as f:
        yolo_obj_time = json.load(f)
    time_diff = time.time() - yolo_obj_time["time"]
    if time_diff > 1:
        time.sleep(0.1)
        if time_diff > 10:
            print("yolo_json time out. perhaps the yolo model is not working")
        return yolo_json(filter_label_ids)
    yolo_obj = yolo_obj_time["objs"]
    if filter_label_ids:
        yolo_obj = [obj for obj in yolo_obj if obj["label_id"] in filter_label_ids]
    return yolo_obj
def recog(searchWith: str, label: Union[int,str]):
    if searchWith == "face" or searchWith == "yolo":
        try:
            label_id=int(label)
        except:
            print(f"label {label} is not a valid int")
            raise ValueError(f"Label {label} is not a valid int. Note that if searchWith is yolo or face, you should write INT label.If searchWith is vllm, you should write STRING label.")
    if searchWith == "face":
        name_id,name,sim,pos = capFaceRecog.main()
        if name_id != label_id:
            return []
        else:
            return [{"label": name, "label_id": name_id, "pos": pos}]
    elif searchWith == "yolo":
        yolo_obj = yolo_json([label_id])
        return yolo_obj
    elif searchWith == "vllm":
        reply = LLMService.funcCallLLM(
            "请识别图片中的物体，并返回物体的名称和像素位置",
            "yolo.png",
            [registerRecogResult],
            []
        )
        if reply.tool_calls:
            tool_call = reply.tool_calls[0]
            if tool_call.function.name == "registerRecogResult":
                return tool_call.function.arguments
    else:
        raise ValueError("Invalid search method. Choose from face, yolo, or vllm")
def findAround(searchWith: str, label: Union[int,str]):
    toDir=0
    armmotion.toDir(toDir,120)
    foundObjs=[]
    while toDir<135:
        print(f"searching in {toDir} degrees")
        resultDict=recog(searchWith,label)
        if resultDict:
            foundObjs.extend(resultDict)
            break
        toDir+=20
        armmotion.toDir(toDir,120)
    toDir=0
    armmotion.toDir(toDir,120)
    while toDir>-135 and len(foundObjs)==0:
        resultDict=recog(searchWith,label)
        if resultDict:
            foundObjs.extend(resultDict)
            break
        toDir-=20
        armmotion.toDir(toDir,120)
    return foundObjs
def FindWalk(searchWith: str, label: Union[int,str], walk: bool, catch: bool, putToPack: bool):
    try:
        foundObjs=findAround(searchWith,label)
        if not foundObjs:
            return {"success":False, "error": "No object found"}
        if walk:
            pose = transformPos.imgPos2BasePos(foundObjs[0])
            if pose:
                wheelmotion.move(pose["x"], pose["y"], 0.15)
        if catch:
            foundObjs=findAround(searchWith,label)
            if foundObjs:
                pose = transformPos.imgPos2BasePos(foundObjs[0])
                if pose:
                    armmotion.release()
                    armmotion.toPos(pose["x"], pose["y"], pose["z"])
                    armmotion.catch()
        if putToPack:
            armmotion.moveToPack()
            armmotion.release()
        return {"success":True}
    except Exception as e:
        print(e, traceback.format_exc())
        return {"success":False, "error":str(e)}
def grabThings(walk: bool):
    # if walk:
    #     wheelmotion.move(0,0,0.15)
    armmotion.grabThings()
    return {"success":True}
def SwitchWheelLeg(wheelOrLeg: int):
    # if wheelOrLeg == 0:
    #     wheelmotion.moveForward(1)
    # else:
    #     legmotion.walk()
    # return {"success":True}
    legmotion.stand()
    legmotion.up()
    wheelmotion.moveForward(1)
def wheelMove(dist:int):
    wheelmotion.moveForward(1)
    time.sleep(1)
    wheelmotion.moveForward(0)
transMap = {
    "ForwardModel": "前进",
    "TurnModel": "转向",
    "FindWalkModel": "寻找并行走",
    "CaptureModel": "拍照",
    "SpeakModel": "说话",
    "SetScheduleModel": "设置日程",
    "SwitchWheelLegModel": "切换轮腿模式",
    "WheelMoveModel": "轮子移动"
}
def router(name:str, args:dict, taskagent):
    if name not in ["SpeakModel","grabThingsModel"]:
        Speak(f"正在{transMap.get(name,name)}")
    if name == "ForwardModel":
        return Forward(args["distance"])
    elif name == "TurnModel":
        return Turn(args["direction"])
    elif name == "ArmMoveModel":
        return ArmMove(args["x"], args["y"], args["z"], args["releaseOnPos"], args["catchOnPos"])
    elif name == "CaptureModel":
        return Capture(args["yaw_degree"], args["pitch_degree"])
    elif name == "SpeakModel":
        return Speak(args["content"])
    elif name == "SetScheduleModel":
        return SetSchedule(args["jobtype"], args["task"], args["hour"], args["minute"], args["second"], taskagent)
    elif name == "FindWalkModel":
        return FindWalk(args["searchWith"], args["label"], args["walk"], args["catch"], args["putToPack"])
    elif name == "grabThingsModel":
        return grabThings(args["walk"])
    elif name == "SwitchWheelLegModel":
        return SwitchWheelLeg(args["wheelOrLeg"])
    elif name== "WheelMoveModel":
        return wheelMove(args["dist"])
    else:
        return {"error": "Invalid tool name"}
if __name__ == "__main__":
    # print(FindWalk("yolo",83,True,True,True))
    SwitchWheelLeg(0)
# def job():
#     print("Running scheduled task...")
# def job2():
#     print("Running scheduled task2...")
# scheduler = BackgroundScheduler()
# scheduler.add_job(job, 'cron', hour=10, minute=11)
# scheduler.add_job(job2, 'interval', seconds=10)
# scheduler.start()



